package com.zhao.algorithm.code


import com.typesafe.config.Config
import org.apache.kafka.clients.consumer.{ConsumerConfig, ConsumerRecord}
import org.apache.kafka.common.TopicPartition
import org.apache.kafka.common.serialization.StringDeserializer
import org.apache.spark.streaming.StreamingContext
import org.apache.spark.streaming.dstream.InputDStream
import org.apache.spark.streaming.kafka010.{ConsumerStrategies, KafkaUtils, LocationStrategies}

/**
 * Description: <br/>
 * Copyright (c) ，2020 ， 赵 <br/>
 * This program is protected by copyright laws. <br/>
 * Date： 2020/11/23 11:43
 *
 * @author 柒柒
 * @version : 1.0
 */

object CustomKafkaConsumer {
  def getKafkaTopicDS(ssc: StreamingContext, load: Config, offsets: Map[TopicPartition, Long]) = {

    val bootstrapServers = load.getString("algorithm.kafka.consumer.bootstrapServers")
    val topicName = load.getString("algorithm.kafka.consumer.topic")
    val groupId = load.getString("algorithm.kafka.consumer.groupId")
    val offset = load.getString("algorithm.kafka.consumer.offset")
    val maxPoll = load.getInt("algorithm.kafka.consumer.maxPoll")
    val isAutoCommit = false: java.lang.Boolean

    println("init kafka consumer ok : ", topicName, groupId, bootstrapServers, offset, maxPoll, isAutoCommit)

    val kafkaParams = Map(
      ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG -> bootstrapServers,
      ConsumerConfig.GROUP_ID_CONFIG -> groupId,
      ConsumerConfig.AUTO_OFFSET_RESET_CONFIG -> offset,
      ConsumerConfig.MAX_POLL_RECORDS_CONFIG -> maxPoll.toString,
      ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG -> isAutoCommit,
      ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG -> classOf[StringDeserializer],
      ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG -> classOf[StringDeserializer]
    )

    // 判断offsets是否为空
    var kafkaTopicDS: InputDStream[ConsumerRecord[String, String]] = null

    if (offsets.isEmpty) {
      println("offset获取是空，走默认的创建")
      kafkaTopicDS = KafkaUtils.createDirectStream(
        ssc,
        LocationStrategies.PreferConsistent,
        ConsumerStrategies.Subscribe[String, String](Set(topicName), kafkaParams)
      )
    } else {
      println("offset不是空的，不走默认的创建")
      kafkaTopicDS = KafkaUtils.createDirectStream(
        ssc,
        LocationStrategies.PreferConsistent,
        ConsumerStrategies.Subscribe[String, String](Set(topicName), kafkaParams, offsets)
      )
    }


    kafkaTopicDS
  }
}
